import databricks.koalas as ks
import numpy as np
import pandas as pd
import pyspark.pandas as ps

if __name__ == '__main__':
    kdf = ps.DataFrame({'A': ['foo', 'bar', 'foo', 'bar', 'foo', 'bar', 'foo', 'foo'],
                        'B': ['one', 'one', 'two', 'three', 'two', 'two', 'one', 'three'],
                        'C': np.random.randn(8),
                        'D': np.random.randn(8)})
    kdf1 = kdf.groupby('A').sum()
    print(kdf1)
    kdf2 = kdf.groupby(['A', 'B']).sum()
    print(kdf2)
